import os
import shutil
import re
import argparse
from pathlib import Path

def organize_single_cell_files(input_dir, output_dir, file_extensions=None):
    """
    组织单细胞分析文件的Python脚本
    
    参数:
    input_dir: 输入目录路径
    output_dir: 输出目录路径
    file_extensions: 可能的文件扩展名列表
    """
    
    if file_extensions is None:
        file_extensions = ['.tsv', '.tsv.gz', '.mtx', '.mtx.gz', '.csv', '.csv.gz']
    
    # 创建输出目录
    Path(output_dir).mkdir(parents=True, exist_ok=True)
    
    # 获取所有相关文件
    all_files = []
    for ext in file_extensions:
        all_files.extend(Path(input_dir).glob(f"*{ext}"))
    
    print(f"找到 {len(all_files)} 个相关文件")
    
    # 识别样本和文件类型
    sample_files = {}
    
    for file_path in all_files:
        filename = file_path.name
        
        # 尝试识别样本名称 - 这里需要根据您的文件命名约定调整
        sample_name = identify_sample_name(filename)
        
        # 识别文件类型
        file_type = identify_file_type(filename)
        
        if sample_name and file_type:
            if sample_name not in sample_files:
                sample_files[sample_name] = {}
            
            sample_files[sample_name][file_type] = file_path
            print(f"样本 '{sample_name}' - {file_type}: {filename}")
    
    # 为每个样本创建文件夹并复制文件
    processed_samples = 0
    
    for sample_name, files in sample_files.items():
        # 检查是否包含所有三种必需文件
        required_files = ['barcodes', 'features', 'matrix']
        if all(req_file in files for req_file in required_files):
            # 创建样本目录
            sample_dir = Path(output_dir) / sample_name
            sample_dir.mkdir(exist_ok=True)
            
            # 标准文件名映射
            standard_names = {
                'barcodes': '',
                'features': '', 
                'matrix': ''
            }
            
            # 复制并重命名文件
            for file_type, original_path in files.items():
                # 保持原始压缩状态
                original_ext = ''.join(original_path.suffixes)
                target_name = standard_names[file_type].replace('.gz', original_ext if original_ext.endswith('.gz') else '')
                
                target_path = sample_dir / target_name
                shutil.copy2(original_path, target_path)
                print(f"  复制: {original_path.name} -> {target_name}")
            
            processed_samples += 1
            print(f"✓ 样本 '{sample_name}' 处理完成")
        else:
            missing_files = set(required_files) - set(files.keys())
            print(f"✗ 样本 '{sample_name}' 缺少文件: {missing_files}")
    
    print(f"\n处理完成! 成功处理 {processed_samples} 个样本")

def identify_sample_name(filename):
    """
    根据文件名识别样本名称
    这里需要根据您的文件命名约定进行调整
    """
    # 示例模式 - 请根据您的实际文件命名调整这些正则表达式
    
    # 模式1: 样本名在文件名开头，后面跟着文件类型描述
    patterns = [
        r'^([A-Za-z0-9_-]+?)[._](?:barcodes|features|matrix|genes|counts)',
        r'^([A-Za-z0-9_-]+?)[._](?:tsv|mtx|csv|gz)',
        r'^(Sample_[A-Za-z0-9_-]+)',
        r'^([A-Za-z0-9]+_[A-Za-z0-9]+)'
    ]
    
    for pattern in patterns:
        match = re.search(pattern, filename, re.IGNORECASE)
        if match:
            return match.group(1)
    
    # 如果无法识别，返回文件名的主要部分
    return filename.split('.')[0]

def identify_file_type(filename):
    """
    识别文件类型：barcodes, features, 或 matrix
    """
    filename_lower = filename.lower()
    
    # 识别barcodes文件
    if any(keyword in filename_lower for keyword in ['barcode', 'barcodes']):
        return 'barcodes'
    
    # 识别features文件 (可能是features.tsv或genes.tsv)
    elif any(keyword in filename_lower for keyword in ['feature', 'features', 'gene', 'genes']):
        return 'features'
    
    # 识别matrix文件
    elif any(keyword in filename_lower for keyword in ['matrix', 'mtx', 'count', 'counts']):
        return 'matrix'
    
    return None

def main():
    parser = argparse.ArgumentParser(description='组织单细胞分析文件到样本文件夹')
    parser.add_argument('--input', '-i', required=True, help='输入目录路径')
    parser.add_argument('--output', '-o', required=True, help='输出目录路径')
    
    args = parser.parse_args()
    
    print("单细胞分析文件组织工具")
    print("=" * 50)
    print(f"输入目录: {args.input}")
    print(f"输出目录: {args.output}")
    print("=" * 50)
    
    organize_single_cell_files(args.input, args.output)

if __name__ == "__main__":
    # 如果直接运行脚本，使用示例
    if len(os.sys.argv) == 1:
        print("使用方法示例:")
        print("python script.py --input /path/to/raw/files --output /path/to/organized/files")
        print("\n或者直接修改下面的路径运行:")
        
        # 在这里设置您的路径
        input_directory = "/path/to/your/raw/files"  # 修改为您的输入路径
        output_directory = "/path/to/your/organized/files"  # 修改为您的输出路径
        
        if input_directory != "/path/to/your/raw/files":
            organize_single_cell_files(input_directory, output_directory)
    else:
        main()